package site.edody.dframe.websocket;

import org.java_websocket.WebSocket;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * DFrame WebSocket 会话缓存
 * @author Dody
 */
public class DFrameWebSocketSessionCache {
    private static final Logger log = LoggerFactory.getLogger(DFrameWebSocketSessionCache.class);

    // 存储结构: Map<serverName, Map<clientId, WebSocket>>
    private final ConcurrentHashMap<String, ConcurrentHashMap<String, WebSocket>> sessionCache =
            new ConcurrentHashMap<>();

    public void addSession(String serverName, String clientId, WebSocket webSocket) {
        sessionCache.computeIfAbsent(serverName, k -> new ConcurrentHashMap<>())
                .put(clientId, webSocket);
        log.debug("[DFrame-{}] Added session: {}, total sessions: {}",
                serverName, clientId, getSessionCount(serverName));
    }

    public void removeSession(String serverName, String clientId) {
        ConcurrentHashMap<String, WebSocket> serverSessions = sessionCache.get(serverName);
        if (serverSessions != null) {
            serverSessions.remove(clientId);
            log.debug("[DFrame-{}] Removed session: {}, remaining sessions: {}",
                    serverName, clientId, getSessionCount(serverName));
        }
    }

    public WebSocket getSession(String serverName, String clientId) {
        ConcurrentHashMap<String, WebSocket> serverSessions = sessionCache.get(serverName);
        return serverSessions != null ? serverSessions.get(clientId) : null;
    }

    public Map<String, WebSocket> getServerSessions(String serverName) {
        ConcurrentHashMap<String, WebSocket> serverSessions = sessionCache.get(serverName);
        return serverSessions != null ? new ConcurrentHashMap<>(serverSessions) : new ConcurrentHashMap<>();
    }

    public String findSessionId(String serverName, WebSocket webSocket) {
        ConcurrentHashMap<String, WebSocket> serverSessions = sessionCache.get(serverName);
        if (serverSessions != null) {
            return serverSessions.entrySet().stream()
                    .filter(entry -> entry.getValue().equals(webSocket))
                    .map(Map.Entry::getKey)
                    .findFirst()
                    .orElse(null);
        }
        return null;
    }

    public int getSessionCount(String serverName) {
        ConcurrentHashMap<String, WebSocket> serverSessions = sessionCache.get(serverName);
        return serverSessions != null ? serverSessions.size() : 0;
    }

    public void clearServerSessions(String serverName) {
        sessionCache.remove(serverName);
        log.info("[DFrame-{}] Cleared all sessions", serverName);
    }

    public Map<String, Integer> getAllSessionCounts() {
        Map<String, Integer> counts = new ConcurrentHashMap<>();
        sessionCache.forEach((serverName, sessions) ->
                counts.put(serverName, sessions.size()));
        return counts;
    }
}